feat(multi-gpu): mark the remaining text encoders as idle-GPU offloadable by Pfannkuchensack · Pull Request #9428 · invoke-ai/InvokeAI

Pfannkuchensack · 2026-08-01T01:48:38Z

Summary

The idle_gpu_offloadable flag introduced in #9311 lets an encoder-only node run its whole execution on a borrowed idle GPU when offload_text_encoders_to_idle_gpus is enabled. Every text encoder that existed at the time was marked, but four have been added since and were never given the flag — so on a multi-GPU machine they always run on the session's own GPU, holding VRAM that the generation needs.

Marks the four stragglers: wan_text_encoder, krea2_text_encoder, ideogram4_text_encoder, ernie_image_text_encoder.

All four satisfy the condition the flag documents ("encoder-only nodes that store their result on the CPU and do no work on the session's own GPU"): each returns via detach().to("cpu") and persists only the conditioning name. Krea-2's optional mask input is a TensorField that is passed through untouched and first processed in krea2_denoise.

After this change 14 text-encoder invocations carry the flag (15 including flux_redux, which is an image-prompt encoder rather than a text one).

Tests

The flag lives on the @invocation decorator, so nothing inside a node body hints that the CPU-store contract exists — a future edit dropping a .to("cpu") would break multi-GPU silently. Added the regression tests that flux2_klein and flux_redux already have:

a registry-driven guard that every *_text_encoder node carries the flag (enumerating rather than listing the four, so the next encoder added is what it catches), with a documented _NOT_OFFLOADABLE escape hatch;
per-node tests that each of the four detaches and moves its conditioning to the CPU, plus one asserting Krea-2 forwards its regional-mask TensorField untouched — resolving it in the encoder would pull a tensor onto the borrowed GPU.

Borrow-cost documentation

device_pool.py claimed a lending session "waits out the (short) encoder node". That was already inaccurate: the borrow spans the node's model load, and caches are per-device, so the first borrow of a GPU always cold-loads the encoder there. That cost amortizes — later borrows are sticky and hit the cache — but work that recurs per execution does not.

ernie_image_text_encoder is the one node here where that matters: its optional prompt enhancer runs an autoregressive generate() of up to 1024 tokens inside the borrow, so it re-stalls the lent GPU on every generation instead of once. Corrected the claim in device_pool.py, the idle_gpu_offloadable docstring, and the node-authoring guide so the tradeoff is visible.

Reviewers may want to weigh in on ERNIE specifically. Marking it is still a net win in the common case and it is kept as-is here, but the options are not obvious — see the review comment below for the analysis, including why a per-instance opt-out (offload only when the enhancer is off) is worse than either alternative given how buildErnieImageGraph emits its positive/negative encoder pair.

Related Issues / Discussions

Follows up #9311 (which introduced the flag) and #9263 (multi-GPU parallel session execution).

QA Instructions

Requires two or more CUDA GPUs and offload_text_encoders_to_idle_gpus: true in invokeai.yaml.

Run a generation with Wan, Krea-2, Ideogram4 or ERNIE-Image while a second GPU is idle.
With InvokeAI at DEBUG level, the session processor logs Running <node type> on idle device cuda:N (session device cuda:M) for the text-encoder node. Before this PR that line never appeared for these four.
Confirm the conditioning still produces identical images — the node runs on a different device but its output is moved to the CPU either way. (Not applicable to ERNIE-Image with the prompt enhancer on: the rewrite is sampled, so it is not reproducible run to run, with or without this PR.)
Single-GPU installs are unaffected: with no idle device to borrow, the node runs exactly as before.

Merge Plan

Nothing special. No node versions are bumped because no fields changed:
idle_gpu_offloadable is a ClassVar set by the @invocation decorator, not a
pydantic field, so it does not appear in the node schema and schema.ts is
untouched. Bumping the version here would be actively harmful — it signals a
template change to the frontend and causes saved workflows to re-instantiate the
node, for no benefit.

Checklist

The PR has a short but descriptive title, suitable for a changelog
Tests added / updated (if applicable)
❗Changes to a redux slice have a corresponding migration — n/a, backend only
Documentation added / updated (if applicable) — borrow-cost criterion added to the node-authoring guide
Updated What's New copy (if doing a release after this PR)

@invocation

…able The `idle_gpu_offloadable` flag (invoke-ai#9311) lets an encoder-only node run on a borrowed idle GPU when `offload_text_encoders_to_idle_gpus` is enabled. Every text encoder that existed at the time was marked; the four added since were not, so on a multi-GPU machine they always occupy the session's own GPU. Marks wan, krea2, ideogram4 and ernie_image. All four meet the flag's stated condition: each returns its result via `detach().to("cpu")` and saves only the conditioning name, doing no work on the session device. Krea-2's optional `mask` input is passed through as a TensorField and is not processed until the denoise node. No schema change: the flag is a ClassVar set by the @invocation decorator, not a pydantic field, so no node version bumps are needed.

@invocation

…ders The four encoders this PR marks store their conditioning on the CPU, which is what makes the borrowed-GPU handoff safe -- but the flag lives on the @invocation decorator, so nothing in the node bodies hints that the contract exists. Add the regression tests that already exist for flux2_klein and flux_redux: - a registry-driven guard that every *_text_encoder node carries the flag, so the next encoder to be added is caught rather than the four already fixed; - per-node tests that each of the four detaches and moves its conditioning to the CPU, plus that krea2 forwards its regional mask TensorField untouched (resolving it here would pull a tensor onto the borrowed GPU). Also correct the borrow-cost documentation. device_pool.py claimed a lending session "waits out the (short) encoder node"; the borrow actually spans the node's model load, and caches are per-device, so the first borrow of a GPU always cold-loads the encoder there. That cost amortizes across later borrows, which hit the cache -- but work that recurs per execution does not amortize. ernie_image_text_encoder runs an autoregressive generate() for its optional prompt enhancer inside the borrow, so it stalls the lent GPU on every generation rather than once. Noted for reviewers in the flag docs and the node-authoring guide. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

lstein · 2026-08-01T19:08:30Z

Reviewed this and pushed one commit to the branch (ab47529338) — tests plus a documentation correction. No behavioural change: all four markings are as you had them. Details below, including one design question I'd like your read on.

The markings check out

I traced each of the four against the contract. All four store their conditioning via detach().to("cpu") and persist only the conditioning name; Krea-2's regional mask is a TensorField name reference that is forwarded untouched and first resolved in krea2_denoise. Attacks that failed, for the record:

Device-dependent output dtype. krea2_text_encoder.py:157 stamps the conditioning dtype from choose_bfloat16_safe_dtype(borrowed_device), which would be wrong on a heterogeneous pool — but krea2_denoise.py:159 re-casts to the session device's inference_dtype, so it cannot mismatch.
Loaders resolving the wrong GPU. All four loaders derive target_device from TorchDevice.choose_torch_device(), which returns the thread-local pin first, so models land on the borrowed GPU as intended.
Borrow lifecycle. The finally in _maybe_offload_to_idle_gpu restores the pin, the cache-stats swap and the lock on both the exception and CanceledException paths; the borrow wraps only invoke_internal, so output serialization is outside it.
Missed candidates. Every *_text_encoder node is now marked. (wan_ref_image_encoder is a VAE image encoder that does image I/O and calls TorchDevice.empty_cache() — correctly out of scope here, but perhaps worth its own look.)

Two small notes: the body said "all 12 text-encoder invocations" — it's 14 after this PR (15 marked including flux_redux, which isn't a text encoder); and QA step 3 isn't checkable for ERNIE with the enhancer on, since the rewrite is sampled. Both corrected in the description.

What I added

The flag lives on the decorator, so nothing in a node body signals that the CPU-store contract exists — dropping a .to("cpu") later would break multi-GPU silently with nothing failing. flux2_klein and flux_redux each already have a regression test for exactly this; the four new ones now do too. The marker test enumerates the registry rather than listing the four, so what it actually guards is the next encoder somebody adds — which is the failure mode this PR is fixing.

The one thing I'd like your read on: ERNIE

device_pool.py said a lending session "waits out the (short) encoder node". That was already inaccurate before this PR — the borrow spans the node's model load, and caches are per-device, so the first borrow of a GPU always cold-loads the encoder there. That amortizes (borrow selection is sticky, later borrows hit the cache). What doesn't amortize is work that recurs per execution, and ernie_image_text_encoder runs the prompt enhancer's autoregressive generate() — up to 1024 tokens, not interruptible by the cancel event — inside the borrow. So it re-stalls the lent GPU on every generation: a session dequeued onto that GPU logs Executing queue item N and then blocks in acquire_session() for the length of someone else's prompt rewrite.

I tried fixing that with a per-instance opt-out (offload only when the enhancer won't run) and backed it out — it's worse than either alternative. buildErnieImageGraph.ts:54-73,99-102 emits two ernie_image_text_encoder nodes sharing one text_encoder: pos_prompt with the enhancer on, and neg_prompt with it always off (guidance_scale > 1, the default). A per-instance rule splits that pair across GPUs: the positive node runs on the session GPU, loading the encoder and the PE there — exactly what the feature exists to avoid — and the negative node then borrows the idle GPU and cold-loads the same encoder a second time, for no benefit, since it's already resident. Strictly worse than either marking ERNIE or not marking it.

So the real options are:

Keep it marked (what's on the branch). Both encoder nodes borrow the same idle GPU, one encoder copy, session GPU free — the full benefit. Cost: the idle GPU's lock is held through the PE on every generation. Only bites when a second session is dequeued during that window, so on a single-user multi-GPU box it costs nothing.
Don't mark ERNIE. No stall, no benefit; the encoder and PE stay on the session GPU contending with the ERNIE transformer.
Split the prompt enhancer into its own node, leaving a genuinely encoder-only ernie_image_text_encoder. That gets both properties, but it's a node-schema change and a graph-builder change — a separate PR.

I left (1) in place since it's what you wrote and it's defensible, and documented the tradeoff in device_pool.py, the flag docstring, and the node-authoring guide so the next person marking a node weighs runtime rather than just "is it an encoder". Happy to go to (2) if you'd rather be conservative, and (3) seems like the right eventual answer.

lstein

LGTM and am happy to merge when you give the word. However, please take a look at the comment regarding the handling of ERNIE Image and let me know if you want to make the suggested architectural change of splitting the PE into its own node.

`ernie_image_text_encoder` is `idle_gpu_offloadable`, so its whole execution runs on a borrowed idle GPU whose exclusive-use lock is held until the node returns. An encoder forward is short and its model load amortizes into the borrowed device's cache; the bundled prompt enhancer's autoregressive `generate()` — up to 1024 tokens — runs afresh on every generation and never amortizes, so it re-stalled the lent GPU each time. A session dequeued onto that GPU logged `Executing queue item N` and then blocked in `acquire_session()` for the length of someone else's prompt rewrite. Carve the rewrite out into `ernie_image_prompt_enhancer`, a StringOutput node that is deliberately *not* offloadable. The encoder becomes genuinely encoder-only and keeps the flag; the enhancer stays on the session's own GPU. A per-instance opt-out was not viable: `buildErnieImageGraph` emits two encoders sharing one text encoder, and gating on the enhancer would split that pair across GPUs and cold-load the same encoder twice. The graph builder now wires the enhancer between the prompt node and the positive encoder when the toggle is on, and wires the prompt straight through otherwise. With no enhancer connected the node passes the prompt through, so pipelines that ship no PE submodel behave as before. Also make the rewrite cancelable: `generate()` gets a StoppingCriteria bound to the session's cancel event, and a cancelled run raises rather than encoding the truncated prompt. `ernie_image_text_encoder` goes to 2.0.0 — the six enhancer fields are removed, which breaks saved workflows that set them. It is a Prototype node. Tests pin the split in both directions: the enhancer must not become offloadable, and the encoder must not regain enhancer fields (either change would silently invalidate its flag with nothing else failing). Plus passthrough, the token cap, cancellation, and the new graph wiring in both toggle states.

Main's invoke-ai#9428 marked every text-encoder node idle_gpu_offloadable and added a registry guard asserting the marker on all *_text_encoder nodes; the merge brought that guard onto this branch where flux2_dev_text_encoder (which neither parent knew about) fails it. The flag alone would be wrong: the marker's contract is that the saved conditioning is CPU-backed, because the borrowed GPU's pool lock is released the moment the node returns. Move the Mistral embeds to CPU before save (the placeholder clip_embeds follows their device), add the marker, bump to 1.0.1, and add the same output-device regression test the Klein encoder has. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* feat(flux2): add FLUX.2 [dev] support Adds end-to-end support for FLUX.2 [dev] alongside the existing Klein implementation. Dev uses Mistral Small 3.1 (24B) as its sole text encoder instead of Klein's Qwen3, with joint_attention_dim=15360 and the guidance-distilled 32B transformer. Backend - taxonomy: Flux2VariantType.Dev, ModelType.MistralEncoder, ModelFormat.MistralEncoder, MistralVariantType - configs: probe dev via context_in_dim=15360 (main + LoRA); new mistral_encoder.py with Diffusers / Checkpoint / GGUF configs; Main_Diffusers_Flux2_Config accepts Flux2Pipeline class name - loaders: new mistral_encoder.py (AutoModel for Diffusers folder, MistralModel for single-file + GGUF with llama.cpp key conversion). Existing Klein transformer loaders are generic enough for dev - ModelRecordChanges.variant union extended with MistralVariantType Invocations - flux2_dev_model_loader, flux2_dev_text_encoder (Mistral chat-template with FLUX2_DEV_SYSTEM_MESSAGE and layer-stacking 10/20/30), flux2_dev_lora_loader (+ collection variant) - MistralEncoderField on model.py; flux2_denoise / flux2_vae_decode / flux2_vae_encode reused unchanged (already model-agnostic) Frontend - types/hooks/selectors for MistralEncoder, isFlux2DevMainModelConfig, selectFlux2DevDiffusersModels, useMistralEncoderModels - params slice fields flux2DevVaeModel / flux2DevMistralEncoderModel / flux2DevSourceModel + reducers, selectIsFlux2Dev / selectIsFlux2Klein - ParamFlux2DevModelSelect component, wired into AdvancedSettingsAccordion - buildFLUXGraph dev branch with full txt2img / img2img / inpaint / outpaint + multi-reference image editing (same flux_kontext + collect chain as Klein, since Flux2RefImageExtension is model-agnostic) - addFlux2DevLoRAs helper for dev LoRA wiring - zModelType / zModelFormat / zFlux2VariantType extended for mistral_encoder / mistral_small_3_1 / dev - OpenAPI schema regenerated, TS types updated Starter models - FLUX.2 [dev] Diffusers (bf16 + NF4), three GGUFs (Q4/Q6/Q8), Mistral encoder (bf16 + NF4) * fix(flux2): wire dev path end-to-end, harden Mistral encoder loader Follow-up fixes after first end-to-end run with FLUX.2 [dev] GGUF + Mistral 3.x GGUF + standalone FLUX.2 VAE. Frontend - buildFLUXGraph: wire dev model loader's vae into both flux2_denoise (required for BN statistics / inpaint) and flux2_vae_decode; missing edge was raising RequiredConnectionException at runtime - readiness.ts: variant-aware FLUX.2 readiness check — dev requires flux2DevVaeModel + flux2DevMistralEncoderModel (or a Dev diffusers source); Klein keeps Qwen3/VAE check. Threads hasFlux2DevDiffusersSource through generate + canvas tabs and updates buildGenerateTabArg / buildCanvasTabArg test helpers - en.json: noFlux2DevVaeModelSelected, noFlux2DevMistralEncoderModelSelected Mistral encoder loader (GGUF / single-file) - Fix "Cannot copy out of meta tensor": llama.cpp conversion produced `model.*` keys but loader instantiated bare MistralModel (no `model.` prefix). Add _convert_for_bare_mistral_model to strip the prefix and drop lm_head before load_state_dict - _materialize_remaining_meta_tensors: after load_state_dict, replace any still-meta parameters (norms→ones, others→zeros) and buffers so the cache→VRAM move can't fail on partial state dicts, with a warning listing what was missing - llama.cpp converter: map attn_q_norm/attn_k_norm (Mistral 3.x qk-norm variants), with ordering before attn_q/attn_k to avoid bad rewrites Tokenizer / processor fallback - _load_processor_with_offline_fallback walks a list of sources (black-forest-labs/FLUX.2-dev tokenizer subfolder, then mistralai/Mistral-Small-3.1-… and 3.2-…), trying AutoProcessor then AutoTokenizer for each, cache-first then online. Final error spells out the three workarounds (install Diffusers folder, set HF_ENDPOINT, pre-cache the tokenizer) - flux2_dev_text_encoder: try multimodal `[{type, text}]` chat template first (PixtralProcessor / Mistral3Processor), fall back to plain string content (AutoTokenizer), then to manual [INST]…[/INST] Qwen3 encoder probe strictness - _get_qwen3_variant_from_state_dict and _get_variant_from_config now return None / raise NotAMatchError for unknown hidden_size instead of silently defaulting to qwen3_4b. The old fallback meant any llama.cpp GGUF causal LM (Mistral, Llama, …) was wrongly classified as Qwen3 — visible when the Mistral 3.x GGUF was identified as a Qwen3-4B encoder - Checkpoint / GGUF / Diffusers loaders propagate the strictness * Chore Path fix * FLUX.2 [dev]: restrict Mistral encoder to 30-layer cow + add recall handlers Upstream Mistral Small 3.1/3.2 (40 layers) produces off-distribution embeddings under FLUX.2's static (10, 20, 30) hidden-state extraction. The joint attention was actually trained against BFL's 30-layer cow-mistral3-small distillation — both Comfy-Org's safetensors and gguf-org's cow GGUFs ship the same 30-layer weights, just packaged differently. - Probing (configs/mistral_encoder.py) now rejects non-cow Mistrals across all three formats (Diffusers / Checkpoint / GGUF) with a clear error. - Loader (load/model_loaders/mistral_encoder.py) extracts the embedded Tekken tokenizer from the `tekken_model` U8 (safetensors) / fp16-per-byte (cow GGUF) tensor via mistral_common, falling back to the BFL HF tokenizer. Removes the INVOKEAI_MISTRAL_TOKENIZER_SOURCE env var. - Starter models: drop upstream Mistral 3.x entries, add Comfy-Org bf16/fp8/fp4 variants alongside the cow GGUFs. - MistralVariantType: drop Small3_1, keep only Cow. - pyproject.toml: add mistral-common dependency. Frontend recall: - Add Flux2DevVAEModel + Flux2DevMistralEncoderModel handlers, disambiguating Klein vs dev via presence of `mistral_encoder` / `qwen3_encoder` metadata fields (both bases are `flux2`). - Wire both into the Recall Parameters panel (hardcoded list was missing them). - Add `metadata.mistralEncoder` i18n key + colocated tests. * feat(flux2-dev): match ComfyUI's Mistral reference + accept 40-layer encoders After studying ComfyUI's `Flux2Tokenizer` / `Mistral3_24BModel` reference implementation, align the FLUX.2 [dev] text-encoder path with their setup: - Probing now accepts both 30-layer (cow distillation) and 40-layer (Mistral Small 3, BFL canonical / upstream) Mistrals. Re-adds `MistralVariantType.Mistral24B` alongside `Cow`. All three configs (Diffusers / Checkpoint / GGUF) updated. - Loaders strip `model.norm` (replace with Identity) when the loaded weights are the 30-layer cow distillation. Matches Comfy's `final_norm=False` for the pruned variant; for transformers' `MistralModel` the final RMSNorm is always built but the cow was trained against the raw post-layer-29 state. - 40-layer loads now log a clear warning that upstream Mistral 3.1 / 3.2 is NOT what FLUX.2's joint attention was trained against and recommends the Comfy-Org bf16/fp8/fp4 or gguf-org cow GGUF variants. BFL's canonical bundled text_encoder is also 40-layer so we don't hard-reject; the warning is opt-in self-discipline. - Text encoder invocation switches from `apply_chat_template(messages, ...)` to a raw text template `[SYSTEM_PROMPT]{sys}[/SYSTEM_PROMPT][INST]{prompt}[/INST]` fed straight to the tokenizer — byte-for-byte matches Comfy's `Flux2Tokenizer.llama_template.format(text)`. System prompt now includes the literal `\n` between "object" and "attribution" Comfy ships. - `_TekkenChatTemplateAdapter` renamed to `_TekkenRawTextAdapter` and exposes a `__call__(text, padding_side='left', ...)` interface that Tekken-encodes the raw string (BOS=1, no EOS) and left-pads with token id 11. Matches Comfy's `pad_left=True` / `pad_token=11` settings. Frontend types extended for the new `mistral3_24b` variant (zMistralVariantType, MODEL_VARIANT_TO_LONG_NAME, schema.ts). * fix(ui): remove unused exports flagged by knip on FLUX.2 [dev] branch Knip reported 6 unused exports. Each was dead code rather than incomplete wiring, verified against the actual consumers: - Drop the vestigial `flux2DevSourceModel` param end-to-end (state field, default, migration, reducer, action, selector, test). The FLUX graph builder auto-picks the diffusers source itself and never read this param; no UI set it. Mirrors how the Klein path already works. - Delete `selectIsFlux2Klein`; the graph builder computes this locally and only `selectIsFlux2Dev` is consumed. - Un-export `zMistralVariantType`; used only in the local `zAnyModelVariant` union, like `zQwenImageVariantType`. - Delete `selectMistralEncoderModels`; components use the `useMistralEncoderModels` hook instead. - Un-export `isFlux2DevMainModelConfig`; used only within types.ts, like its `isFluxDevMainModelConfig` / `isFlux2Klein9BMainModelConfig` siblings. * Chore OpenApi * Chore Ruff * chore(deps): lock mistral-common for FLUX.2 [dev] Mistral encoder * fix(flux2): disambiguate dev/Klein VAE recall by model variant The dev-vs-Klein VAE recall keyed off the presence of a mistral_encoder metadata field, but that field is only written when a standalone Mistral encoder is selected. A FLUX.2 [dev] image whose encoder came from a Diffusers source has a vae field but no mistral_encoder, so its VAE was silently recalled into the Klein slice. Resolve the image's own main model and check variant === 'dev' instead — the same signal the graph builder uses. Add regression coverage for the mistral_encoder-absent dev case, and add the missing modelManager.flux2Dev* i18n keys so the [dev] VAE/encoder labels are translatable. * fix(flux2): pass prompt as text= keyword to Mistral processor The diffusers FLUX.2-dev text encoder loads a PixtralProcessor, whose first positional __call__ parameter is `images`, not `text`. Passing the prompt positionally routed it into `images`, breaking the diffusers encoder path (only single-file/GGUF encoders, which use a text-first adapter, had been exercised). Pass text= explicitly. * fix(flux2): pass prompt as text= keyword to Mistral processor The diffusers FLUX.2-dev text encoder loads a PixtralProcessor, whose first positional __call__ parameter is `images`, not `text`. Passing the prompt positionally raised "Incorrect image source", breaking the diffusers encoder path entirely. Only single-file/GGUF encoders (text-first adapter) had been exercised. Verified against transformers 5.5.4. fix(flux2): emit Tekken special tokens in the embedded-tokenizer adapter _TekkenRawTextAdapter used mistral_common's raw Tekkenizer.encode, which runs with SpecialTokenPolicy.IGNORE and BPE-encodes the FLUX.2 markers ([SYSTEM_PROMPT], [/SYSTEM_PROMPT], [INST], [/INST]) as literal text — 54 tokens instead of 36, corrupting the prompt structure fed to FLUX.2 on the single-file and GGUF paths. Resolve the marker ids from the tokenizer's special vocab and splice them in; output is now byte-identical to the reference PixtralProcessor. * Add FLux2.dev to readme * fix(flux2-dev): address review — regional guidance, model classification, LoRA guards, encoder probes - Wire FLUX.2 [dev] regional guidance through addRegions instead of dropping it silently - Require pipeline layout for Main_Diffusers_Flux2_Config so transformer-only checkouts don't register as broken main models - Reject Klein<->dev LoRA cross-wiring on both frontend (variant filter) and backend (loaders raise) - Discriminate non-Mistral GGUFs via vocab-size floor; accept text_encoder.-prefixed encoder layouts at probe time - Dequantize fp8 checkpoints per-tensor to target dtype and drop lm_head before casting (avoid whole-dict fp32 peak) - Raise on unexpected Mistral layer count instead of inventing extraction indices - Fail Klein VAE recall closed when the main model is unresolvable - Add missing modelManager.mistralEncoder i18n key - Dedup: single-pass GGUF metadata read, consistent norm materialization, cat-based conditioning, drop redundant t() defaultValues * Feat: FLUX.2 [dev] review fixes, dedup, and shared-source refactors Address the PR #9234 review (correctness, install-probe gaps, polish) plus the deduplication follow-ups. Correctness - Wire FLUX.2 [dev] regional guidance through addRegions instead of silently dropping it (posCondCollect + flux2_dev_text_encoder handling case) - Require a full pipeline layout (model_index.json / transformer/) for Main_Diffusers_Flux2_Config so transformer-only checkouts don't register as broken main models and OSError mid-queue - Reject Klein<->dev LoRA cross-wiring on both ends: frontend filters LoRAs by variant in both graph builders; dev/Klein loaders raise instead of warn - Discriminate non-Mistral GGUFs via a vocab-size floor so Llama-2-13B and similar 5120-hidden/40-layer LMs no longer install as Mistral encoders - Accept text_encoder.-prefixed encoder layouts at install probe (matches the loader's prefix stripping) - Dequantize fp8 Mistral checkpoints per-tensor to the target dtype and drop lm_head before casting (avoid a whole-dict fp32 transient that can OOM) - Raise on an unexpected Mistral layer count instead of inventing extraction indices that silently degrade output - Fail Klein VAE recall closed when the image's main model is unresolvable - Add the missing modelManager.mistralEncoder i18n key Dedup / single source of truth - Consolidate the FLUX.2 dimension->variant tables (context/vec/hidden) into a shared configs/flux2_variant.py used by main.py and lora.py - Mistral loaders key the final-RMSNorm / warning decision on config.variant instead of re-deriving from num_hidden_layers==30 - Merge the separate Klein/dev VAE redux slots into one flux2VaeModel (slice migration v3->v4) and collapse the two metadata VAE handlers into one, removing the recall-disambiguation - Parameterize the near-identical dev/Klein canvas graph blocks into one shared addFlux2Features closure; add dev-path coverage to buildFLUXGraph.test.ts - Single-pass GGUF metadata read, consistent norm materialization, cat-based conditioning tensor, and drop redundant t() defaultValues in ParamFlux2DevModelSelect Tests: model_identification suite green; frontend parsing / graph / readiness / modelSelected suites green. * Fix: bump paramsSlice persist version to 4 for the shared FLUX.2 VAE slot The v3->v4 migration (Klein/dev VAE slots -> flux2VaeModel) bumped _version but left zParamsState._version at literal(3) and the initial state at 3, so migrate()'s final zParamsState.parse rejected with "expected 3". Bump the schema literal + initial state to 4 and add a v3->v4 migration test. * Fix: address FLUX.2 [dev] round-2 review (4 blockers + 6 cleanups) Blockers: - params migration: seed flux2DevMistralEncoderModel in the v3->v4 step so a genuine v3 blob passes zParamsState.parse() instead of wiping the whole params slice on upgrade; rebuild the migration test fixture as a field-accurate v3 object so it actually covers the regression. - guidance for [dev]: resolve the image's own model in the Guidance metadata parse gate and exempt variant === 'dev' so guidance is displayed/recalled for [dev] (still skipped for Klein); render the guidance slider for FLUX.2 [dev]. - source-model variant guard: require variant == Dev where the dev loader validates its Mistral/VAE source, and reject a [dev] source in the Klein loader — a mismatched pipeline otherwise fails with an opaque matmul error in denoise. - tokenizer offline load: drop the dead root-dir fallback + duplicated pre-try and add a root-directory AutoProcessor step to _load_tokenizer_for_model so processor files alongside the encoder weights load offline. Cleanups: - extract _reinit_inv_freq() with a rope_theta -> rope_parameters/rope_scaling fallback (fixes a latent AttributeError on pinned transformers 5.5, removes a verbatim duplicate). - flux2_dev_lora_collection_loader: replace the base assert with a ValueError that rejects non-FLUX.2 LoRAs, mirroring the Klein collection loader. - diffusers Mistral load: drop the never-run vision_tower/multi_modal_projector (~0.8GB) so they stay out of the cache and VRAM transfers. - clear flux2DevMistralEncoderModel on base switch and intra-flux2 variant switch. - pin mistral-common>=1.5.4,<2 (validated against 1.11.6). - fix contradictory 40-layer docstrings to match the taxonomy/loader story. * Chore openapi * fix(ui): bump params persist schema to v5 to resolve the dual-v4 collision main and this branch both shipped _version 4 with different new keys (PiD fields vs the flux2 VAE merge + Mistral encoder slot), so a v4 blob written by either parent would fail zParamsState.parse() after the merge and wipe the whole params slice. Keep main's v3->v4 step verbatim and move the flux2 slot merge + Mistral seed to a new v4->v5 step with conditional seeding for both v4 shapes. Also seed the five Wan component fields in v3->v4: they were added to the schema without a version bump while releases were still writing v3 blobs, so a genuine released-build (v6.13.x) v3 blob fails parse() on them today - same wipe, inherited from main. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Chore openapi * fix(flux2): scope the cross-variant source guard to encoder extraction, widen the tokenizer ladder The FLUX.2 loaders ran one validator on both the VAE- and the encoder-extraction call site, so the cross-variant check also rejected VAE-only sourcing. Klein and [dev] share the same 32-channel AutoencoderKLFlux2 and the linear UI relies on that -- buildFLUXGraph falls back to any FLUX.2 diffusers pipeline when only the VAE is needed, and readiness does not filter by variant. A Klein GGUF main plus a standalone Qwen3 encoder plus a [dev] pipeline as the only diffusers model therefore hit a ValueError behind an enabled Invoke button. Split the validator: format-only for the VAE path, format + variant for the encoder path. The Mistral tokenizer ladder's local-directory rungs tried AutoProcessor only. On transformers 5.5.4 that raises OSError for a mistral3 config.json without preprocessor_config.json -- exactly the BFL-style standalone-encoder layout the rungs target -- so the ladder fell through to the HF fetch and failed offline. Both rungs now loop (AutoProcessor, AutoTokenizer), with KeyError in the except tuple for tekken-only directories. * Chore openapi * feat(metadata): declare mistral_encoder on core_metadata, bump to 2.2.0 FLUX.2 [dev] recorded its Mistral text encoder as an undeclared extra key, relying on the node's `extra='allow'`, while the Klein counterpart `qwen3_encoder` is a proper field. Declare it so it lands in the OpenAPI schema and is typed in the frontend instead of `unknown`. Also bumps the node version, which was left at 2.1.0 across several model integrations that widened the node: `ideogram4_caption` (#9303) and the generation modes for FLUX.2, Anima, Qwen-Image, Ideogram 4, Wan, Krea-2 and Ernie. All changes are additive, so this is a minor bump - saved workflows carrying a core_metadata node now auto-update to the current field set on load rather than silently keeping a stale one. * fix(flux2): make flux2_dev_text_encoder idle-GPU-offloadable Main's #9428 marked every text-encoder node idle_gpu_offloadable and added a registry guard asserting the marker on all *_text_encoder nodes; the merge brought that guard onto this branch where flux2_dev_text_encoder (which neither parent knew about) fails it. The flag alone would be wrong: the marker's contract is that the saved conditioning is CPU-backed, because the borrowed GPU's pool lock is released the moment the node returns. Move the Mistral embeds to CPU before save (the placeholder clip_embeds follows their device), add the marker, bump to 1.0.1, and add the same output-device regression test the Klein encoder has. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(flux2): stop the Mistral tokenizer ladder from crashing and from silently mis-encoding Tekken Round-4 review blockers, both reproduced against transformers 5.5.4 with real files before fixing. 1. An AttributeError from a probe rung killed the whole load. `AutoTokenizer.from_pretrained` on a directory whose `tokenizer_config.json` names a `tokenizer_class` the installed transformers does not know resolves that class to None and dereferences it without a guard — exactly the layout `MistralCommonTokenizer.save_pretrained` writes. That is not in `_TOKENIZER_LOAD_ERRORS`, so it escaped `_try_load_tokenizer_from_dir` and crashed a load the HF rung would have completed. The probes now catch broadly; the expected-error tuple only selects the log level, so an unexpected failure is still logged loudly instead of being swallowed. 2. A root `tekken.json` next to `config.json` did not fail in `AutoTokenizer` — it resolved to a mistral-common-backed tokenizer that BPE-encodes `[SYSTEM_PROMPT]`/`[INST]` as literal text instead of splicing them as single Tekken ids. The encode "worked" and conditioning was silently off-distribution. Fixed on two independent paths: the ladder now reads a standalone `tekken.json` itself, ahead of the transformers probes, and any mistral-common-backed result is re-wrapped in `_TekkenRawTextAdapter` through its underlying `MistralTokenizer` rather than used as-is. The vocab is fine — only its `__call__` is wrong — so re-wrapping beats discarding, which would have traded silent corruption for an offline RuntimeError. Verified: the re-wrapped ids are identical to the reference adapter's. Also closes both non-blockers: - `_validate_encoder_source` in the Klein loader rejected only [dev], so a Klein 9B pipeline passed as `qwen3_source_model` for a Klein 4B transformer and hit the very matmul error the guard exists to prevent (the frontend and the standalone-encoder path both enforce the family match; the workflow editor's source field was the only way in). It is now an allowlist keyed on a shared `_KLEIN_TO_QWEN3_VARIANT` map — mirroring the frontend's `KLEIN_TO_QWEN3_VARIANT_MAP` — and checks the source's Qwen3 family against the main model, so a future third FLUX.2 variant fails closed on the Klein side too, not just on [dev]. `_validate_qwen3_encoder_variant` shares that map and now uses `getattr` instead of `hasattr`, which turned a None variant into an AttributeError in the error path rather than the intended ValueError. - The [dev] loader's `_validate_diffusers_format` docstring claimed the linear UI relies on the permissive VAE path. That holds for Klein, but the [dev] builder sources from dev-only pipelines and readiness gates on one, so there the cross-variant VAE case is reachable through the workflow editor only. The justification now states what actually holds: the 32-channel AutoencoderKLFlux2 is shared (the repo ships the Klein-sourced `flux2_vae` as a dependency of every [dev] GGUF starter), and `mistral_source_model` is not variant-filtered in the editor. Tests: a structurally valid Tekken fixture, so the ladder exercises the success path rather than only the raise path the previous fake produced; regression tests for both blockers on the directory and HF rungs; Klein-family coverage including same-family acceptance and the standalone-encoder guard's negative path, which had no coverage at all. All new tests mutation-verified — reverting the broad catch, the tekken rung, the re-wrap, the family check, or the allowlist each fails at least one. tests/app + tests/backend/model_manager: 3010 passed. The 9 failures are the pre-existing network-dependent ones in test_model_install / test_load_api / test_download_queue. --------- Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Pfannkuchensack requested review from JPPhoto, blessedcoolant, dunkeroni and lstein as code owners August 1, 2026 01:48

github-actions Bot added python PRs that change python files invocations PRs that change invocations labels Aug 1, 2026

Merge branch 'main' into feat/idle_gpu_offloadable_encoders

a0bea75

lstein self-assigned this Aug 1, 2026

github-actions Bot added backend PRs that change backend files python-tests PRs that change python tests docs PRs that change docs labels Aug 1, 2026

lstein approved these changes Aug 1, 2026

View reviewed changes

lstein and others added 3 commits August 2, 2026 11:58

Merge branch 'main' into feat/idle_gpu_offloadable_encoders

59f7c17

Merge branch 'main' into feat/idle_gpu_offloadable_encoders

53cf745

github-actions Bot added the frontend PRs that change frontend files label Aug 3, 2026

Pfannkuchensack requested a review from lstein August 3, 2026 00:27

lstein enabled auto-merge (squash) August 3, 2026 00:37

Merge branch 'main' into feat/idle_gpu_offloadable_encoders

578d5f5

lstein merged commit 9f2b2be into invoke-ai:main Aug 3, 2026
17 checks passed

Pfannkuchensack deleted the feat/idle_gpu_offloadable_encoders branch August 3, 2026 01:16

lstein mentioned this pull request Aug 5, 2026

Feat: flux2 dev support #9234

Merged

5 tasks

Navigation Menu

Uh oh!

feat(multi-gpu): mark the remaining text encoders as idle-GPU offloadable - #9428

feat(multi-gpu): mark the remaining text encoders as idle-GPU offloadable#9428
lstein merged 7 commits into
invoke-ai:mainfrom
Pfannkuchensack:feat/idle_gpu_offloadable_encoders

Pfannkuchensack commented Aug 1, 2026 •

edited by lstein

Loading

Uh oh!

lstein commented Aug 1, 2026

Uh oh!

lstein left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

Pfannkuchensack commented Aug 1, 2026 • edited by lstein Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Tests

Borrow-cost documentation

Related Issues / Discussions

QA Instructions

Merge Plan

Checklist

Uh oh!

lstein commented Aug 1, 2026

The markings check out

What I added

The one thing I'd like your read on: ERNIE

Uh oh!

lstein left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Pfannkuchensack commented Aug 1, 2026 •

edited by lstein

Loading